-- Copyright (C) 2016 Koz Ross <koz.ross@retro-freedom.nz>

-- This program is free software: you can redistribute it and/or modify
-- it under the terms of the GNU General Public License as published by
-- the Free Software Foundation, either version 3 of the License, or
-- (at your option) any later version.

-- This program is distributed in the hope that it will be useful,
-- but WITHOUT ANY WARRANTY; without even the implied warranty of
-- MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
-- GNU General Public License for more details.

-- You should have received a copy of the GNU General Public License
-- along with this program.  If not, see <http://www.gnu.org/licenses/>.

module Data.Logic.DNF (Atom,
                       Clause,
                       Form,
                       Interpretation,
                       combine,
                       interpret, interpretSafe, invert,
                       minSatisfying, mkAtom, mkClause, mkForm, mkInterpretation,
                       optimize) where

import BasicPrelude hiding (empty)

import qualified Data.Set as S (deleteAt, deleteFindMin,
                                elemAt, empty,
                                filter, findMin, fromList,
                                insert,
                                map, member,
                                null,
                                singleton, size,
                                toAscList, toList)

import qualified Data.IntMap as M (fromList, insert, lookup)

import Test.QuickCheck

import Data.List.Extra (nubOn)

data Atom = Neg { index :: Int }
          | Pos { index :: Int }
          deriving (Eq, Ord, Show, Read)

instance Arbitrary Atom where
  arbitrary = sized $ \i -> do
    b <- arbitrary
    bool Neg Pos b <$> choose (0, abs i)

mkAtom :: Bool -> Int -> Maybe Atom
mkAtom b ix = bool Neg Pos b ix <$ guard (ix >= 0)

newtype Clause = Clause { atoms :: Set Atom }
  deriving (Eq, Show, Read) 

instance Ord Clause where
  compare = comparing (Down . S.size . atoms) <> comparing atoms

instance Arbitrary Clause where
  arbitrary = sized $ \i -> do
    j <- choose (1, max 1 (abs i))
    Clause . S.fromList . take j . nubOn index <$> infiniteList

-- does the uncommented thing have the same semantics as the commented things?
mkClause :: [Atom] -> Maybe Clause
mkClause x = (Clause . S.fromList $ x) <$ guard (not . null $ x)

newtype Form = Form { clauses :: Set Clause }
  deriving (Eq, Ord, Show, Read)

instance Arbitrary Form where
  arbitrary = sized $ \i -> do
    j <- choose (1, max 1 (abs i))
    Form . S.fromList . take j . nub <$> infiniteList

mkForm :: [Clause] -> Maybe Form
mkForm x = (Form . S.fromList $ x) <$ guard (not . null $ x)

combine :: Form -> Form -> Form
combine f1 f2 = Form (clauses f1 ++ clauses f2)

newtype Interpretation = Interp {
  toMap :: IntMap Bool }
  deriving (Eq, Ord, Show, Read)  

instance Arbitrary Interpretation where
  arbitrary = sized $ \i -> do
    limit <- choose (0, abs i)
    Interp . M.fromList . zip [0 .. limit] <$> infiniteList

mkInterpretation :: [Bool] -> Interpretation
mkInterpretation = Interp . M.fromList . zip [0 ..]

minSatisfying :: Clause -> Interpretation
minSatisfying = Interp . M.fromList . fmap toPair . S.toList . atoms
  where toPair (Pos n) = (n, True)
        toPair (Neg n) = (n, False)

-- A safe, lazy interpreter
-- Will try and find an answer even in the presence of Nothings if it makes
-- logical sense
-- Will Nothing-out if forced to do so
-- E.g. a v (b ^ c) with interpretation a = True, b = True:
--    1) a -> True by interp
--    2) (b ^ c) -> Nothing
--    3) a v (b ^ c) -> Just True
-- However, a v (b ^ c) with interpretation a = False, b = True:
--    1) a -> False by interp
--    2) (b ^ c) -> Nothing
--    3) a v (b ^ c) -> Nothing
interpret :: Interpretation -> Form -> Maybe Bool
interpret = evalForm'

-- Backup interpreter - guaranteed safe, but slow
interpretSafe :: Interpretation -> Form -> Maybe Bool
interpretSafe = evalForm

-- Functions to support the interpreters

{-# INLINE evalAtom #-}
evalAtom :: Interpretation -> Atom -> Maybe Bool
evalAtom (Interp i) (Pos x) = M.lookup x i
evalAtom (Interp i) (Neg x) = not <$> M.lookup x i

evalClause :: Interpretation -> Clause -> Maybe Bool
evalClause i (Clause x)
  | S.member (Just False) evaluated = Just False
  | evaluated == S.singleton (Just True) = Just True
  | otherwise = Nothing
  where evaluated = S.map (evalAtom i) x

evalClause' :: Interpretation -> Clause -> Maybe Bool
evalClause' i (Clause x) = foldr1 f . fmap (evalAtom i) $ S.toList x
  where f b@(Just False) _ = b
        f (Just True) y = y
        f Nothing y@(Just False) = y
        f Nothing _ = Nothing

evalForm :: Interpretation -> Form -> Maybe Bool
evalForm i (Form x)
  | S.member (Just True) evaluated = Just True
  | evaluated == S.singleton (Just False) = Just False
  | otherwise = Nothing
  where evaluated = S.map (evalClause i) x

evalForm' :: Interpretation -> Form -> Maybe Bool
evalForm' i (Form x) = foldr1 f . fmap (evalClause' i) $ S.toList x
  where f b@(Just True) _ = b
        f (Just False) y = y
        f Nothing y@(Just True) = y
        f Nothing _ = Nothing

-- A simple optimizer
-- Nothings-out on a contradiction
optimize :: Form -> Maybe Form
optimize x = reduceForm . deimplicate <$> dedup x

-- A negator
-- Can get ugly on large forms
-- Does not optimize!
invert :: Form -> Form
invert = Form . S.fromList . fmap Clause . foldl1' extendAll . toIntermediate

-- Functions to support the negator

-- Assumes second list is of singletons!
extendAll :: [Set Atom] -> [Set Atom] -> [Set Atom]
extendAll olds news = catMaybes (extend <$> olds <*> news)

-- Assumes second set is a singleton!
extend :: Set Atom -> Set Atom -> Maybe (Set Atom)
extend old new = let only = S.findMin new in
  if inByIndex only old
  then case only of
    Pos i -> if S.member (Neg i) old
             then Nothing
             else Just . S.insert only $ old
    Neg i -> if S.member (Pos i) old
             then Nothing
             else Just . S.insert only $ old
  else Just . S.insert only $ old

inByIndex :: Atom -> Set Atom -> Bool
inByIndex x s = S.member (index x) (S.map index s)

toIntermediate :: Form -> [[Set Atom]]
toIntermediate = S.toAscList . S.map inner . clauses
  where inner = S.toAscList . S.map (S.singleton . opposite) . atoms

opposite :: Atom -> Atom
opposite (Pos i) = Neg i
opposite (Neg i) = Pos i

-- Functions to support the optimizer

reduceForm :: Form -> Form
reduceForm f = Form . S.map (reduceClause f) . clauses $ f 

reduceClause :: Form -> Clause -> Clause
reduceClause f (Clause c) = go S.empty (variants c)
  where go res [] = Clause res
        go res ((e, rs) : xs) = let base = minSatisfying . Clause $ rs
                                    exI b = Interp . M.insert (index e) b . toMap
                                    positive = exI True base
                                    negative = exI False base in
          if (interpret positive f == Just True)
             && (interpret negative f == Just True)
          then go res xs
          else go (S.insert e res) xs

variants :: Set a -> [(a, Set a)]
variants s = (\i -> (S.elemAt i s, S.deleteAt i s)) <$> [0 .. (S.size s - 1)]

dedup :: Form -> Maybe Form
dedup f = let dedupped = S.filter uniq . clauses $ f in
  Form dedupped <$ guard (not . S.null $ dedupped)

uniq :: Clause -> Bool
uniq (Clause x) = (S.size . S.map index $ x) == S.size x

deimplicate :: Form -> Form
deimplicate = Form . go S.empty . clauses
  where go acc curr
          | S.null curr = acc
          | otherwise = let (x, xs) = S.deleteFindMin curr in
              if S.null xs
              then go (S.insert x acc) xs
              else let leftHand = minSatisfying x in
                if interpret leftHand (Form xs) == Just True
                then go acc xs
                else go (S.insert x acc) xs
